from typing import Union
from .ollama_service import OllamaService

class GraphQueryService:
    def __init__(self):
        self.ollama_service = OllamaService()
        self.keyword_extraction_prompt = """
        请从下列问题中提取最关键的专业术语或关键词，确保用于在中文故障图谱中检索。
        例如输入
        "为什么压缩压力不足"，应输出
        "压缩压力不足"。

        如果无法识别出专业关键词，请回答“无法提取关键词”。

        问题：
        {question}
        关键词：
        """

    async def extract_keywords(self, question: str) -> Union[str, None]:
        """
        从用户问题中提取关键词
        :param question: 用户问题
        :return: 提取的关键词或None
        """
        if not question or not question.strip():
            print("关键词提取失败：问题为空")
            return None

        try:
            # 构建提示词
            prompt = self.keyword_extraction_prompt.format(question=question)

            # 流式生成关键词
            keywords = ""
            async for chunk in self.ollama_service.generate_stream(prompt):
                keywords += chunk

            # 清理和验证关键词
            return self._clean_and_validate_keywords(keywords)

        except Exception as e:
            print(f"关键词提取失败: {str(e)}")
            return None

    def _clean_and_validate_keywords(self, raw_keywords: str) -> Union[str, None]:
        """
        清理和验证提取的关键词
        """
        # 去除首尾空白和特殊字符
        cleaned = raw_keywords.strip().replace('"', '').replace("'", "").replace("\n", "")

        # 检查无效响应
        if not cleaned:
            return None
        if "无法提取关键词" in cleaned:
            return None
        if "error" in cleaned.lower():
            return None

        # 提取有效关键词（最多3个）
        keywords_list = [kw.strip() for kw in cleaned.split(",") if kw.strip()]
        return ",".join(keywords_list[:3])




